from train import solver,solver2
import torch
from MyNet2.mynet2 import MyNet2
from PoolNet.poolnet import build_model, weights_init

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

'''
SINet_model = SINet(64)
SINet_model.load_state_dict(torch.load("./model_pth/29.pth"))
SI_solver = solver(SINet_model, device, epoch=30, lr=1e-4)

# SI_solver.train("./model_pth/")

SI_solver.test("./output/")
'''

mynet = MyNet2()


# poolnet
PoolNet = build_model('resnet')
PoolNet.cuda()
PoolNet.eval()
PoolNet.apply(weights_init)
PoolNet.base.load_pretrained_model(torch.load('pretrained/resnet18.pth'))

# mynet.load_state_dict(torch.load("./model_pth/my/29.pth"))
mynet_solver = solver2(PoolNet, device, epoch=30, lr=1e-4, root="../data/ECSSD/")
mynet_solver.train("./model_pth/my/")
mynet_solver.test("./output/my/", epoch=29)
